# -*- coding: utf-8 -*-
# @Time : 2024-11-2024/11/30
# @File : visualization.py

import cv2
from utils.image_utils import cv2_add_chinese_text


def draw_skeleton(img, keypoints, connections, joint_angles, angles=None, thickness=3, show_keypoint_idx=True):
    """
    在图像上绘制骨架
    
    参数:
    - img: 目标图像
    - keypoints: 关键点坐标 (x, y)
    - connections: 骨架连接列表，每项为(索引1, 索引2)
    - joint_angles: 关节角度配置
    - angles: 计算出的角度值字典
    - thickness: 线的粗细
    - show_keypoint_idx: 是否显示关键点索引
    
    返回:
    - 带有骨架的图像
    """
    # 颜色定义
    colors = {
        'limb': (255, 255, 255),  # 白色骨架
        'joint': (0, 255, 255)    # 黄色关节点
    }
    
    h, w = img.shape[:2]
    
    # 绘制连接线
    for connection in connections:
        idx1, idx2 = connection
        if idx1 < len(keypoints) and idx2 < len(keypoints):
            pt1 = tuple(map(int, keypoints[idx1]))
            pt2 = tuple(map(int, keypoints[idx2]))
            if pt1[0] > 0 and pt1[1] > 0 and pt2[0] > 0 and pt2[1] > 0:  # 确保点有效且在图像范围内
                try:
                    cv2.line(img, pt1, pt2, colors['limb'], thickness)
                except Exception as e:
                    print(f"绘制线条错误: {e}, 点1: {pt1}, 点2: {pt2}")
    
    # 绘制关节点
    for i, point in enumerate(keypoints):
        try:
            x, y = int(point[0]), int(point[1])
            if x > 0 and y > 0 and x < w and y < h:  # 确保点在图像范围内
                cv2.circle(img, (x, y), thickness + 2, colors['joint'], -1)
                # 添加关键点索引
                if show_keypoint_idx:
                    cv2.putText(img, str(i), (x+5, y+5), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 255, 0), 1)
        except Exception as e:
            print(f"绘制关节点错误: {e}, 点: {point}")
    
    # 如果有计算出的角度，显示它们
    if angles:
        # 按类别分组角度
        angle_categories = {}
        for angle_info in joint_angles:
            name = angle_info["name"]
            if name in angles:
                category = get_angle_category(name)
                if category not in angle_categories:
                    angle_categories[category] = []
                angle_categories[category].append((name, angles[name], angle_info["color"]))
        
        # 在图像右侧显示分组角度
        y_offset = 40
        for category, angle_list in angle_categories.items():
            # 添加类别标题
            img = cv2_add_chinese_text(img, f"【{category}】", (w - 180, y_offset), font_size=16, color=(255, 255, 255))
            y_offset += 20
            
            # 添加该类别下的角度值
            for name, value, color in angle_list:
                text = f"{name}: {value:.1f}°"
                img = cv2_add_chinese_text(img, text, (w - 180, y_offset), font_size=15, color=color)
                y_offset += 18
            
            y_offset += 5  # 类别间隔
        
        # 在关节点位置显示角度值
        for angle_info in joint_angles:
            name = angle_info["name"]
            if name in angles:
                try:
                    p2_idx = angle_info["points"][1]  # 角度顶点的索引
                    if p2_idx < len(keypoints):
                        pos = tuple(map(int, keypoints[p2_idx]))
                        if pos[0] > 0 and pos[1] > 0 and pos[0] < w and pos[1] < h:  # 确保点在图像范围内
                            # 只显示角度值
                            angle_value = angles[name]
                            cv2.putText(img, f"{angle_value:.0f}°", 
                                      (pos[0], pos[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, 
                                      angle_info["color"], 1)
                except Exception as e:
                    print(f"显示角度文本错误: {e}, 关节: {name}")
    
    return img


def get_angle_category(angle_name):
    """根据角度名称获取其类别"""
    upper_limb = ["右肘", "左肘", "右上臂", "左上臂", "手臂摆动", "投掷角度-右"]
    lower_limb = ["右膝", "左膝", "右大腿", "左大腿", "髋膝踝协调-右", "髋膝踝协调-左", "步幅角"]
    trunk = ["躯干前倾", "躯干侧倾", "髋部角度"]
    head = ["头部前倾", "头肩角"]
    
    if angle_name in upper_limb:
        return "上肢"
    elif angle_name in lower_limb:
        return "下肢"
    elif angle_name in trunk:
        return "躯干"
    elif angle_name in head:
        return "头部"
    else:
        return "其他"


def add_frame_info(img, frame_count, additional_info=None, position=(10, 30), font_size=24, color=(0, 255, 0)):
    """添加帧信息到图像"""
    text = f"帧: {frame_count}"
    if additional_info:
        text += f" | {additional_info}"
    
    return cv2_add_chinese_text(img, text, position, font_size, color) 